{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BilmhAzR-4Qg"
      },
      "source": [
        "Copyright 2025 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TsslZVhvNrRy"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pkOtOq0jDE0c"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td>\n",
        "      <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Demos/Emoji-Gemma-on-Web/resources/Convert_Gemma_3_270M_to_ONNX.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cn2YgWGrNv2M"
      },
      "source": [
        "# Convert Gemma 3 270M to ONNX\n",
        "\n",
        "This notebook exports a Gemma 3 model to the ONNX format for use with [Transformers.js](https://huggingface.co/docs/transformers.js/en/index), which uses ONNX Runtime to run models in the browser. The entire process takes under 10 minutes:\n",
        "\n",
        "1. Set up the Colab environment\n",
        "2. Load the model from Hugging Face\n",
        "3. Convert the model with Optimum conversion script\n",
        "4. Test, evaluate, and save the model for further use\n",
        "\n",
        "Gemma 3 270M is designed for task-specific fine-tuning and engineered for efficient performance on mobile, web, and edge devices. You can fine-tune your own model in this [notebook](https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Demos/Emoji-Gemma-on-Web/resources/Fine_tune_Gemma_3_270M_for_emoji_generation.ipynb) and run it in a demo [web app](https://github.com/google-gemini/gemma-cookbook/tree/main/Demos/Emoji-Gemma-on-Web/app-transformersjs) once converted.\n",
        "\n",
        "## Set up development environment\n",
        "\n",
        "The first step is to install packages using pip."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "58TwFUM6SfFR"
      },
      "outputs": [],
      "source": [
        "%pip install transformers==4.56.1 onnx==1.19.0 onnx_ir==0.1.7 onnxruntime==1.22.1 numpy==2.3.2 huggingface_hub"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2g7wTSot_Erp"
      },
      "source": [
        "Restart the session runtime to ensure you're using the newly installed packages."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QuC7NbNy9JXa"
      },
      "source": [
        "## Convert the model\n",
        "To access and save models to Hugging Face, provide your [Access Token](https://huggingface.co/settings/tokens). You can store it as a Colab secret in the left toolbar by specifying `HF_TOKEN` as the 'Name' and adding your unique token as the 'Value'."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6-fGnIinPjdt"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "from huggingface_hub import login\n",
        "hf_token = userdata.get('HF_TOKEN')\n",
        "login(hf_token)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L0PYnb5cODbb"
      },
      "source": [
        "You'll run the build_gemma.py script by [Xenova](https://huggingface.co/Xenova) to convert and quantize the Gemma 3 model.\n",
        "\n",
        "Specify the Hugging Face repo ID of the model to convert. The .onnx exports will be saved to your Colab files."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z292qdELOHOq"
      },
      "outputs": [],
      "source": [
        "!wget https://gist.githubusercontent.com/xenova/a219dbf3c7da7edd5dbb05f92410d7bd/raw/45f4c5a5227c1123efebe1e36d060672ee685a8e/build_gemma.py\n",
        "\n",
        "model_author = \"\"                                         #@param {type:\"string\"}\n",
        "gemma_model = \"myemoji-gemma-3-270m-it\"                   #@param {type:\"string\"}\n",
        "\n",
        "repo_id   = f\"{model_author}/{gemma_model}\"               # Model to convert\n",
        "save_path = f\"/content/{gemma_model}-onnx\"                # Path to save resized model\n",
        "\n",
        "!python build_gemma.py \\\n",
        "    --model_name {repo_id} \\\n",
        "    --output {save_path} \\\n",
        "    -p fp32 fp16 q4 q4f16\n",
        "\n",
        "print(f\"Converted ONNX models saved to {save_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40ae5e9e"
      },
      "source": [
        "## Test the converted model\n",
        "\n",
        "After the .onnx models have saved to your Colab session, try testing inference using ONNX Runtime. Note that this may differ from the ONNX Runtime Web version used for browser inference.\n",
        "\n",
        "Experiment with different text inputs in `text_to_translate`  and explore the performance of different quantized versions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hz8E9RorWWVn"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoConfig, AutoTokenizer, GenerationConfig\n",
        "import onnxruntime\n",
        "import numpy as np\n",
        "\n",
        "# Load config, processor, and model\n",
        "config = AutoConfig.from_pretrained(save_path)\n",
        "generation_config = GenerationConfig.from_pretrained(save_path)\n",
        "tokenizer = AutoTokenizer.from_pretrained(save_path)\n",
        "\n",
        "model_file = \"onnx/model.onnx\"          #@param [\"onnx/model.onnx\", \"onnx/model_fp16.onnx\", \"onnx/model_q4.onnx\", \"onnx/model_q4f16.onnx\"]\n",
        "\n",
        "model_path = f\"{save_path}/{model_file}\"\n",
        "decoder_session = onnxruntime.InferenceSession(model_path)\n",
        "\n",
        "## Set config values\n",
        "num_key_value_heads = config.num_key_value_heads\n",
        "head_dim = config.head_dim\n",
        "num_hidden_layers = config.num_hidden_layers\n",
        "eos_token_id = tokenizer.eos_token_id\n",
        "\n",
        "# Prepare inputs\n",
        "text_to_translate = \"i love sushi\"      # @param {type:\"string\"}\n",
        "messages = [\n",
        "  { \"role\": \"system\", \"content\": \"Translate this text to emoji: \" },\n",
        "  { \"role\": \"user\", \"content\": text_to_translate },\n",
        "]\n",
        "\n",
        "inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors=\"np\")\n",
        "input_ids = inputs['input_ids']\n",
        "attention_mask = inputs['attention_mask']\n",
        "batch_size = input_ids.shape[0]\n",
        "past_key_values = {\n",
        "    f'past_key_values.{layer}.{kv}': np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=np.float32)\n",
        "    for layer in range(num_hidden_layers)\n",
        "    for kv in ('key', 'value')\n",
        "}\n",
        "position_ids = np.tile(np.arange(0, input_ids.shape[-1]), (batch_size, 1))\n",
        "\n",
        "# 3. Generation loop\n",
        "max_new_tokens = 8\n",
        "generated_tokens = np.array([[]], dtype=np.int64)\n",
        "\n",
        "for i in range(max_new_tokens):\n",
        "  logits, *present_key_values = decoder_session.run(None, dict(\n",
        "      input_ids=input_ids,\n",
        "      attention_mask=attention_mask,\n",
        "      position_ids=position_ids,\n",
        "      **past_key_values,\n",
        "  ))\n",
        "\n",
        "  ## Update values for next generation loop\n",
        "  input_ids = logits[:, -1].argmax(-1, keepdims=True)\n",
        "  attention_mask = np.concatenate([attention_mask, np.ones_like(input_ids, dtype=np.int64)], axis=-1)\n",
        "  position_ids = position_ids[:, -1:] + 1\n",
        "\n",
        "  for j, key in enumerate(past_key_values):\n",
        "    past_key_values[key] = present_key_values[j]\n",
        "\n",
        "  generated_tokens = np.concatenate([generated_tokens, input_ids], axis=-1)\n",
        "\n",
        "  if np.isin(input_ids, eos_token_id).any():\n",
        "    break\n",
        "\n",
        "# 4. Output result\n",
        "print(tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7wVBEPJKWSEQ"
      },
      "source": [
        "## Upload to Hugging Face Hub\n",
        "\n",
        "Upload your exported ONNX models to Hugging Face for easy sharing and use."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b4FNJTpKgT7I"
      },
      "outputs": [],
      "source": [
        "import huggingface_hub\n",
        "from huggingface_hub import whoami\n",
        "\n",
        "username = whoami()['name']\n",
        "\n",
        "#@markdown Name your ONNX model:\n",
        "onnx_model = \"myemoji-gemma-3-270m-it-onnx\"       #@param {type:\"string\"}\n",
        "hf_repo_id = f\"{username}/{onnx_model}\"\n",
        "\n",
        "huggingface_hub.create_repo(hf_repo_id, exist_ok=True)\n",
        "\n",
        "repo_url = huggingface_hub.upload_folder(\n",
        "  folder_path=save_path,\n",
        "  repo_id=hf_repo_id,\n",
        "  repo_type=\"model\",\n",
        "  commit_message=f\"Upload ONNX model files for {onnx_model}\"\n",
        "  )\n",
        "\n",
        "print(f\"Uploaded to {repo_url}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cDEBKuXI_V--"
      },
      "source": [
        "## Deploy your model to web with Transformers.js\n",
        "\n",
        "You can now run your Gemma 3 model in the browser using [Transformers.js](https://huggingface.co/docs/transformers.js/en/index) via ONNX Runtime Web. Try it in the [emoji generation web app](https://github.com/google-gemini/gemma-cookbook/tree/main/Demos/Emoji-Gemma-on-Web/app-transformersjs) which runs the model directly in the browser."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Convert_Gemma_3_270M_to_ONNX.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
